package main

import (
	"fmt"
	"log"
	"os"

	"github.com/gin-gonic/gin"
)

type KnowledgeBase struct {
	vectorStore vectorstores.VectorStore
	llm         *openai.LLM
	embedder    embeddings.Embedder
}

// NewKnowledgeBase 创建新的知识库实例
func NewKnowledgeBase(ctx context.Context) (*KnowledgeBase, error) {
	// 初始化LLM
	llm, err := openai.New(
		openai.WithModel("gpt-3.5-turbo"),
		openai.WithToken(os.Getenv("OPENAI_API_KEY")),
	)
	if err != nil {
		return nil, fmt.Errorf("创建LLM失败: %w", err)
	}

	// 初始化Embeddings
	embedder, err := embeddings.NewEmbedder(llm)
	if err != nil {
		return nil, fmt.Errorf("创建Embedder失败: %w", err)
	}

	// 初始化向量数据库 (使用Chroma)
	store, err := chroma.New(
		chroma.WithChromaURL(os.Getenv("CHROMA_URL")),
		chroma.WithEmbedder(embedder),
		chroma.WithDistanceFunction("cosine"),
	)
	if err != nil {
		return nil, fmt.Errorf("创建向量数据库失败: %w", err)
	}

	return &KnowledgeBase{
		vectorStore: store,
		llm:         llm,
		embedder:    embedder,
	}, nil
}

// LoadDocuments 加载本地txt和docx文件到知识库
func (kb *KnowledgeBase) LoadDocuments(ctx context.Context, dirPath string) error {
	// 遍历目录下的所有txt和docx文件
	txtFiles, err := filepath.Glob(filepath.Join(dirPath, "*.txt"))
	if err != nil {
		return fmt.Errorf("读取txt文件失败: %w", err)
	}

	docxFiles, err := filepath.Glob(filepath.Join(dirPath, "*.docx"))
	if err != nil {
		return fmt.Errorf("读取docx文件失败: %w", err)
	}

	// 合并文件列表
	files := append(txtFiles, docxFiles...)

	if len(files) == 0 {
		return fmt.Errorf("未找到任何txt或docx文件")
	}

	var allDocs []schema.Document

	for _, file := range files {
		var docs []schema.Document
		var err error

		// 根据文件扩展名选择合适的加载器
		ext := filepath.Ext(file)
		switch ext {
		case ".txt":
			loader := documentloaders.NewText(file)
			docs, err = loader.Load(ctx)
		case ".docx":
			// 使用docx加载器
			docs, err = kb.loadDocxFile(ctx, file)
		default:
			log.Printf("不支持的文件格式: %s", file)
			continue
		}

		if err != nil {
			log.Printf("加载文件 %s 失败: %v", file, err)
			continue
		}

		allDocs = append(allDocs, docs...)
		log.Printf("成功加载文件: %s, 文档数: %d", file, len(docs))
	}

	if len(allDocs) == 0 {
		return fmt.Errorf("没有成功加载任何文档")
	}

	// 文本分割
	splitter := textsplitter.NewRecursiveCharacter(
		textsplitter.WithChunkSize(500),
		textsplitter.WithChunkOverlap(50),
	)

	splitDocs, err := splitter.SplitDocuments(allDocs)
	if err != nil {
		return fmt.Errorf("文档分割失败: %w", err)
	}

	log.Printf("文档分割完成，共 %d 个片段", len(splitDocs))

	// 将文档添加到向量数据库
	_, err = kb.vectorStore.AddDocuments(ctx, splitDocs)
	if err != nil {
		return fmt.Errorf("添加文档到向量数据库失败: %w", err)
	}

	log.Printf("成功将 %d 个文档片段添加到知识库", len(splitDocs))
	return nil
}

// Query 查询知识库并返回答案
func (kb *KnowledgeBase) Query(ctx context.Context, question string) (string, error) {
	// 相似度搜索
	docs, err := kb.vectorStore.SimilaritySearch(ctx, question, 3)
	if err != nil {
		return "", fmt.Errorf("相似度搜索失败: %w", err)
	}

	if len(docs) == 0 {
		return "抱歉，在知识库中未找到相关信息", nil
	}

	// 构建上下文
	context := ""
	for i, doc := range docs {
		context += fmt.Sprintf("文档片段%d:\n%s\n\n", i+1, doc.PageContent)
	}

	// 构建prompt
	prompt := fmt.Sprintf(`基于以下上下文信息回答问题。如果上下文中没有相关信息，请说"我不知道"。

上下文:
%s

问题: %s

回答:`, context, question)

	// 调用LLM生成答案
	answer, err := kb.llm.Call(ctx, prompt)
	if err != nil {
		return "", fmt.Errorf("LLM调用失败: %w", err)
	}

	return answer, nil
}

// loadDocxFile 加载docx文件
func (kb *KnowledgeBase) loadDocxFile(ctx context.Context, filePath string) ([]schema.Document, error) {
	// 读取docx文件
	r, err := docx.ReadDocxFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("打开docx文件失败: %w", err)
	}
	defer r.Close()

	// 提取文本内容
	docx1 := r.Editable()
	text := docx1.GetContent()

	// 清理文本（移除多余的空白字符）
	text = strings.TrimSpace(text)

	if text == "" {
		return nil, fmt.Errorf("docx文件内容为空")
	}

	// 创建文档对象
	doc := schema.Document{
		PageContent: text,
		Metadata: map[string]any{
			"source":   filePath,
			"filename": filepath.Base(filePath),
			"type":     "docx",
		},
	}

	return []schema.Document{doc}, nil
}

// API请求和响应结构
type QueryRequest struct {
	Question string `json:"question" binding:"required"`
}

type QueryResponse struct {
	Answer string `json:"answer"`
	Error  string `json:"error,omitempty"`
}

func main() {
	ctx := context.Background()

	// 创建知识库实例
	kb, err := NewKnowledgeBase(ctx)
	if err != nil {
		log.Fatalf("初始化知识库失败: %v", err)
	}

	// 加载本地txt文件
	docPath := "./node" // txt文件所在目录
	if err := kb.LoadDocuments(ctx, docPath); err != nil {
		log.Fatalf("加载文档失败: %v", err)
	}

	// 创建Gin路由
	r := gin.Default()

	// 健康检查接口
	r.GET("/health", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"status": "ok",
		})
	})

	// 问答接口
	r.POST("/api/query", func(c *gin.Context) {
		var req QueryRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(400, QueryResponse{
				Error: "请求参数错误: " + err.Error(),
			})
			return
		}

		// 执行查询
		answer, err := kb.Query(ctx, req.Question)
		if err != nil {
			c.JSON(500, QueryResponse{
				Error: "查询失败: " + err.Error(),
			})
			return
		}

		c.JSON(200, QueryResponse{
			Answer: answer,
		})
	})

	// 启动服务
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}

	log.Printf("服务器启动在端口: %s", port)
	if err := r.Run(":" + port); err != nil {
		log.Fatalf("服务器启动失败: %v", err)
	}
}
